import React, { useCallback, useMemo, useState } from "react";
import {
  JsonParam,
  NumberParam,
  StringParam,
  useQueryParam,
} from "use-query-params";
import useLocalStorageState from "use-local-storage-state";
import {
  ColumnPinningState,
  ColumnSort,
  RowSelectionState,
} from "@tanstack/react-table";
import { RotateCw } from "lucide-react";
import findIndex from "lodash/findIndex";
import isObject from "lodash/isObject";
import isNumber from "lodash/isNumber";
import get from "lodash/get";
import {
  useMetricDateRangeWithQueryAndStorage,
  MetricDateRangeSelect,
} from "@/components/pages-shared/traces/MetricDateRangeSelect";

import useTracesOrSpansList, {
  TRACE_DATA_TYPE,
} from "@/hooks/useTracesOrSpansList";
import useTracesOrSpansScoresColumns from "@/hooks/useTracesOrSpansScoresColumns";
import {
  COLUMN_COMMENTS_ID,
  COLUMN_FEEDBACK_SCORES_ID,
  COLUMN_SPAN_FEEDBACK_SCORES_ID,
  COLUMN_GUARDRAIL_STATISTIC_ID,
  COLUMN_GUARDRAILS_ID,
  COLUMN_ID_ID,
  COLUMN_CUSTOM_ID,
  COLUMN_METADATA_ID,
  COLUMN_SELECT_ID,
  COLUMN_TYPE,
  ColumnData,
  ColumnsStatistic,
  DropdownOption,
  DynamicColumn,
  HeaderIconType,
  ROW_HEIGHT,
} from "@/types/shared";
import { CUSTOM_FILTER_VALIDATION_REGEXP } from "@/constants/filters";
import { BaseTraceData, Span, SPAN_TYPE, Trace } from "@/types/traces";
import {
  convertColumnDataToColumn,
  isColumnSortable,
  mapColumnDataFields,
} from "@/lib/table";
import { generateSelectColumDef } from "@/components/shared/DataTable/utils";
import Loader from "@/components/shared/Loader/Loader";
import ExplainerCallout from "@/components/shared/ExplainerCallout/ExplainerCallout";
import NoTracesPage from "@/components/pages/TracesPage/NoTracesPage";
import SearchInput from "@/components/shared/SearchInput/SearchInput";
import FiltersButton from "@/components/shared/FiltersButton/FiltersButton";
import TracesActionsPanel from "@/components/pages/TracesPage/TracesSpansTab/TracesActionsPanel";
import { Separator } from "@/components/ui/separator";
import { Button } from "@/components/ui/button";
import DataTableRowHeightSelector from "@/components/shared/DataTableRowHeightSelector/DataTableRowHeightSelector";
import ColumnsButton from "@/components/shared/ColumnsButton/ColumnsButton";
import DataTable from "@/components/shared/DataTable/DataTable";
import DataTableNoData from "@/components/shared/DataTableNoData/DataTableNoData";
import DataTablePagination from "@/components/shared/DataTablePagination/DataTablePagination";
import LinkCell from "@/components/shared/DataTableCells/LinkCell";
import CodeCell from "@/components/shared/DataTableCells/CodeCell";
import ListCell from "@/components/shared/DataTableCells/ListCell";
import CostCell from "@/components/shared/DataTableCells/CostCell";
import ErrorCell from "@/components/shared/DataTableCells/ErrorCell";
import DurationCell from "@/components/shared/DataTableCells/DurationCell";
import FeedbackScoreCell from "@/components/shared/DataTableCells/FeedbackScoreCell";
import PrettyCell from "@/components/shared/DataTableCells/PrettyCell";
import CommentsCell from "@/components/shared/DataTableCells/CommentsCell";
import FeedbackScoreHeader from "@/components/shared/DataTableHeaders/FeedbackScoreHeader";
import TooltipWrapper from "@/components/shared/TooltipWrapper/TooltipWrapper";
import ThreadDetailsPanel from "@/components/pages-shared/traces/ThreadDetailsPanel/ThreadDetailsPanel";
import TraceDetailsPanel from "@/components/pages-shared/traces/TraceDetailsPanel/TraceDetailsPanel";
import PageBodyStickyContainer from "@/components/layout/PageBodyStickyContainer/PageBodyStickyContainer";
import PageBodyStickyTableWrapper from "@/components/layout/PageBodyStickyTableWrapper/PageBodyStickyTableWrapper";
import TracesOrSpansPathsAutocomplete from "@/components/pages-shared/traces/TracesOrSpansPathsAutocomplete/TracesOrSpansPathsAutocomplete";
import TracesOrSpansFeedbackScoresSelect from "@/components/pages-shared/traces/TracesOrSpansFeedbackScoresSelect/TracesOrSpansFeedbackScoresSelect";
import { formatDate, formatDuration } from "@/lib/date";
import useTracesOrSpansStatistic from "@/hooks/useTracesOrSpansStatistic";
import { useDynamicColumnsCache } from "@/hooks/useDynamicColumnsCache";
import { useIsFeatureEnabled } from "@/components/feature-toggles-provider";
import { FeatureToggleKeys } from "@/types/feature-toggles";
import GuardrailsCell from "@/components/shared/DataTableCells/GuardrailsCell";
import useQueryParamAndLocalStorageState from "@/hooks/useQueryParamAndLocalStorageState";
import { EXPLAINER_ID, EXPLAINERS_MAP } from "@/constants/explainers";
import {
  DetailsActionSection,
  DetailsActionSectionParam,
  DetailsActionSectionValue,
} from "@/components/pages-shared/traces/DetailsActionSection";
import { GuardrailResult } from "@/types/guardrails";
import { SelectItem } from "@/components/ui/select";
import BaseTraceDataTypeIcon from "@/components/pages-shared/traces/TraceDetailsPanel/BaseTraceDataTypeIcon";
import { SPAN_TYPE_LABELS_MAP } from "@/constants/traces";
import SpanTypeCell from "@/components/shared/DataTableCells/SpanTypeCell";
import { Filter } from "@/types/filters";
import {
  USER_FEEDBACK_COLUMN_ID,
  USER_FEEDBACK_NAME,
} from "@/constants/shared";
import { useTruncationEnabled } from "@/components/server-sync-provider";

const getRowId = (d: Trace | Span) => d.id;

const REFETCH_INTERVAL = 30000;

const SPAN_FEEDBACK_SCORE_SUFFIX = " (span)";

/**
 * Formats a score name with the span suffix for display in column labels
 */
const formatSpanScoreLabel = (scoreName: string): string => {
  return `${scoreName}${SPAN_FEEDBACK_SCORE_SUFFIX}`;
};

/**
 * Extracts the score name from a label by removing the span suffix
 */
const parseSpanScoreName = (label: string): string => {
  return label.replace(SPAN_FEEDBACK_SCORE_SUFFIX, "");
};

const SHARED_COLUMNS: ColumnData<BaseTraceData>[] = [
  {
    id: "name",
    label: "Name",
    type: COLUMN_TYPE.string,
  },
  {
    id: "start_time",
    label: "Start time",
    type: COLUMN_TYPE.time,
    accessorFn: (row) => formatDate(row.start_time),
  },
  {
    id: "end_time",
    label: "End time",
    type: COLUMN_TYPE.time,
    accessorFn: (row) => formatDate(row.end_time),
  },
  {
    id: "input",
    label: "Input",
    size: 400,
    type: COLUMN_TYPE.string,
    cell: PrettyCell as never,
    customMeta: {
      fieldType: "input",
    },
  },
  {
    id: "output",
    label: "Output",
    size: 400,
    type: COLUMN_TYPE.string,
    cell: PrettyCell as never,
    customMeta: {
      fieldType: "output",
    },
  },
  {
    id: "duration",
    label: "Duration",
    type: COLUMN_TYPE.duration,
    cell: DurationCell as never,
    statisticDataFormater: formatDuration,
  },
  {
    id: COLUMN_METADATA_ID,
    label: "Metadata",
    type: COLUMN_TYPE.dictionary,
    accessorFn: (row) =>
      isObject(row.metadata)
        ? JSON.stringify(row.metadata, null, 2)
        : row.metadata,
    cell: CodeCell as never,
  },
  {
    id: "tags",
    label: "Tags",
    type: COLUMN_TYPE.list,
    iconType: "tags",
    cell: ListCell as never,
  },
  {
    id: "usage.total_tokens",
    label: "Total tokens",
    type: COLUMN_TYPE.number,
    accessorFn: (row) =>
      row.usage && isNumber(row.usage.total_tokens)
        ? `${row.usage.total_tokens}`
        : "-",
  },
  {
    id: "usage.prompt_tokens",
    label: "Total input tokens",
    type: COLUMN_TYPE.number,
    accessorFn: (row) =>
      row.usage && isNumber(row.usage.prompt_tokens)
        ? `${row.usage.prompt_tokens}`
        : "-",
  },
  {
    id: "usage.completion_tokens",
    label: "Total output tokens",
    type: COLUMN_TYPE.number,
    accessorFn: (row) =>
      row.usage && isNumber(row.usage.completion_tokens)
        ? `${row.usage.completion_tokens}`
        : "-",
  },
  {
    id: "total_estimated_cost",
    label: "Estimated cost",
    type: COLUMN_TYPE.cost,
    cell: CostCell as never,
    explainer: EXPLAINERS_MAP[EXPLAINER_ID.hows_the_cost_estimated],
    size: 160,
  },
];

const DEFAULT_TRACES_COLUMN_PINNING: ColumnPinningState = {
  left: [COLUMN_SELECT_ID, COLUMN_ID_ID],
  right: [],
};

const DEFAULT_TRACES_PAGE_COLUMNS: string[] = [
  "name",
  "input",
  "output",
  "duration",
  COLUMN_COMMENTS_ID,
  USER_FEEDBACK_COLUMN_ID,
];

const SELECTED_COLUMNS_KEY = "traces-selected-columns";
const COLUMNS_WIDTH_KEY = "traces-columns-width";
const COLUMNS_ORDER_KEY = "traces-columns-order";
const COLUMNS_SORT_KEY_SUFFIX = "-columns-sort";
const COLUMNS_SCORES_ORDER_KEY = "traces-scores-columns-order";
const DYNAMIC_COLUMNS_KEY = "traces-dynamic-columns";
const PAGINATION_SIZE_KEY = "traces-pagination-size";
const ROW_HEIGHT_KEY = "traces-row-height";

type TracesSpansTabProps = {
  type: TRACE_DATA_TYPE;
  projectId: string;
  projectName: string;
};

export const TracesSpansTab: React.FC<TracesSpansTabProps> = ({
  type,
  projectId,
  projectName,
}) => {
  const truncationEnabled = useTruncationEnabled();

  const {
    dateRange,
    handleDateRangeChange,
    intervalStart,
    intervalEnd,
    minDate,
    maxDate,
  } = useMetricDateRangeWithQueryAndStorage();
  const [search = "", setSearch] = useQueryParam("search", StringParam, {
    updateType: "replaceIn",
  });

  const [traceId = "", setTraceId] = useQueryParam("trace", StringParam, {
    updateType: "replaceIn",
  });

  const [spanId = "", setSpanId] = useQueryParam("span", StringParam, {
    updateType: "replaceIn",
  });

  const [threadId = "", setThreadId] = useQueryParam("thread", StringParam, {
    updateType: "replaceIn",
  });

  const [page = 1, setPage] = useQueryParam("page", NumberParam, {
    updateType: "replaceIn",
  });

  const [size, setSize] = useQueryParamAndLocalStorageState<
    number | null | undefined
  >({
    localStorageKey: PAGINATION_SIZE_KEY,
    queryKey: "size",
    defaultValue: 100,
    queryParamConfig: NumberParam,
    syncQueryWithLocalStorageOnInit: true,
  });

  const [, setLastSection] = useQueryParam(
    "lastSection",
    DetailsActionSectionParam,
    {
      updateType: "replaceIn",
    },
  );

  const [height, setHeight] = useQueryParamAndLocalStorageState<
    string | null | undefined
  >({
    localStorageKey: ROW_HEIGHT_KEY,
    queryKey: "height",
    defaultValue: ROW_HEIGHT.small,
    queryParamConfig: StringParam,
    syncQueryWithLocalStorageOnInit: true,
  });

  const [filters = [], setFilters] = useQueryParam(
    `${type}_filters`,
    JsonParam,
    {
      updateType: "replaceIn",
    },
  );

  const isGuardrailsEnabled = useIsFeatureEnabled(
    FeatureToggleKeys.GUARDRAILS_ENABLED,
  );
  const [sortedColumns, setSortedColumns] = useQueryParamAndLocalStorageState<
    ColumnSort[]
  >({
    localStorageKey: `${type}${COLUMNS_SORT_KEY_SUFFIX}`,
    queryKey: `${type}_sorting`,
    defaultValue: [],
    queryParamConfig: JsonParam,
  });

  const filtersConfig = useMemo(
    () => ({
      rowsMap: {
        type: {
          keyComponentProps: {
            options: [
              {
                value: SPAN_TYPE.general,
                label: SPAN_TYPE_LABELS_MAP[SPAN_TYPE.general],
              },
              {
                value: SPAN_TYPE.tool,
                label: SPAN_TYPE_LABELS_MAP[SPAN_TYPE.tool],
              },
              {
                value: SPAN_TYPE.llm,
                label: SPAN_TYPE_LABELS_MAP[SPAN_TYPE.llm],
              },
              ...(isGuardrailsEnabled
                ? [
                    {
                      value: SPAN_TYPE.guardrail,
                      label: SPAN_TYPE_LABELS_MAP[SPAN_TYPE.guardrail],
                    },
                  ]
                : []),
            ],
            placeholder: "Select type",
            renderOption: (option: DropdownOption<SPAN_TYPE>) => {
              return (
                <SelectItem
                  key={option.value}
                  value={option.value}
                  withoutCheck
                  wrapperAsChild={true}
                >
                  <div className="flex w-full items-center gap-1.5">
                    <BaseTraceDataTypeIcon type={option.value} />
                    {option.label}
                  </div>
                </SelectItem>
              );
            },
          },
        },
        [COLUMN_METADATA_ID]: {
          keyComponent: TracesOrSpansPathsAutocomplete,
          keyComponentProps: {
            rootKeys: ["metadata"],
            projectId,
            type,
            placeholder: "key",
            excludeRoot: true,
          },
        },
        [COLUMN_CUSTOM_ID]: {
          keyComponent: TracesOrSpansPathsAutocomplete,
          keyComponentProps: {
            rootKeys: ["input", "output"],
            projectId,
            type,
            placeholder: "key",
            excludeRoot: false,
          },
          validateFilter: (filter: Filter) => {
            if (
              filter.key &&
              filter.value &&
              !CUSTOM_FILTER_VALIDATION_REGEXP.test(filter.key)
            ) {
              return `Key is invalid, it should begin with "input", or "output" and follow this format: "input.[PATH]" For example: "input.message" `;
            }
          },
        },
        [COLUMN_FEEDBACK_SCORES_ID]: {
          keyComponent: TracesOrSpansFeedbackScoresSelect,
          keyComponentProps: {
            projectId,
            type,
            placeholder: "Select score",
          },
        },
        ...(type === TRACE_DATA_TYPE.traces
          ? {
              [COLUMN_SPAN_FEEDBACK_SCORES_ID]: {
                keyComponent: TracesOrSpansFeedbackScoresSelect,
                keyComponentProps: {
                  projectId,
                  type: TRACE_DATA_TYPE.spans,
                  placeholder: "Select span score",
                },
              },
            }
          : {}),
        [COLUMN_GUARDRAILS_ID]: {
          keyComponentProps: {
            options: [
              { value: GuardrailResult.FAILED, label: "Failed" },
              { value: GuardrailResult.PASSED, label: "Passed" },
            ],
            placeholder: "Status",
          },
        },
      },
    }),
    [projectId, type, isGuardrailsEnabled],
  );

  const [rowSelection, setRowSelection] = useState<RowSelectionState>({});
  const [isTableDataEnabled, setIsTableDataEnabled] = useState(false);

  // Enable table data loading after initial render to allow users to change the date filter
  React.useEffect(() => {
    const timer = setTimeout(() => {
      setIsTableDataEnabled(true);
    }, 0);
    return () => clearTimeout(timer);
  }, []);

  const { data, isPending, refetch } = useTracesOrSpansList(
    {
      projectId,
      type: type as TRACE_DATA_TYPE,
      sorting: sortedColumns,
      filters,
      page: page as number,
      size: size as number,
      search: search as string,
      truncate: truncationEnabled,
      fromTime: intervalStart,
      toTime: intervalEnd,
    },
    {
      enabled: isTableDataEnabled,
      refetchInterval: REFETCH_INTERVAL,
      refetchOnMount: false,
    },
  );

  const { refetch: refetchExportData } = useTracesOrSpansList(
    {
      projectId,
      type: type as TRACE_DATA_TYPE,
      sorting: sortedColumns,
      filters,
      page: page as number,
      size: size as number,
      search: search as string,
      truncate: false,
      fromTime: intervalStart,
      toTime: intervalEnd,
    },
    {
      enabled: false,
      refetchOnMount: "always",
    },
  );

  const { data: statisticData, refetch: refetchStatistic } =
    useTracesOrSpansStatistic(
      {
        projectId,
        type: type as TRACE_DATA_TYPE,
        filters,
        search: search as string,
        fromTime: intervalStart,
        toTime: intervalEnd,
      },
      {
        refetchInterval: REFETCH_INTERVAL,
      },
    );

  const { data: feedbackScoresData, isPending: isFeedbackScoresPending } =
    useTracesOrSpansScoresColumns(
      {
        projectId,
        type: type as TRACE_DATA_TYPE,
      },
      {
        refetchInterval: REFETCH_INTERVAL,
      },
    );

  const {
    data: spanFeedbackScoresData,
    isPending: isSpanFeedbackScoresPending,
  } = useTracesOrSpansScoresColumns(
    {
      projectId,
      type: TRACE_DATA_TYPE.spans,
    },
    {
      enabled: type === TRACE_DATA_TYPE.traces,
      refetchInterval: REFETCH_INTERVAL,
    },
  );

  const noData = !search && filters.length === 0;
  const noDataText = noData
    ? `There are no ${type === TRACE_DATA_TYPE.traces ? "traces" : "spans"} yet`
    : "No search results";

  const rows: Array<Span | Trace> = useMemo(
    () => data?.content ?? [],
    [data?.content],
  );

  const sortableBy: string[] = useMemo(
    () => data?.sortable_by ?? [],
    [data?.sortable_by],
  );

  const columnsStatistic: ColumnsStatistic = useMemo(
    () => statisticData?.stats ?? [],
    [statisticData],
  );

  const [selectedColumns, setSelectedColumns] = useLocalStorageState<string[]>(
    SELECTED_COLUMNS_KEY,
    {
      defaultValue: DEFAULT_TRACES_PAGE_COLUMNS,
    },
  );

  const [columnsOrder, setColumnsOrder] = useLocalStorageState<string[]>(
    COLUMNS_ORDER_KEY,
    {
      defaultValue: [],
    },
  );

  const [scoresColumnsOrder, setScoresColumnsOrder] = useLocalStorageState<
    string[]
  >(COLUMNS_SCORES_ORDER_KEY, {
    defaultValue: [],
  });

  const [columnsWidth, setColumnsWidth] = useLocalStorageState<
    Record<string, number>
  >(COLUMNS_WIDTH_KEY, {
    defaultValue: {},
  });

  const dynamicScoresColumns = useMemo(() => {
    return (feedbackScoresData?.scores?.slice() ?? [])
      .sort((c1, c2) => c1.name.localeCompare(c2.name))
      .map<DynamicColumn>((c) => ({
        id: `${COLUMN_FEEDBACK_SCORES_ID}.${c.name}`,
        label: c.name,
        columnType: COLUMN_TYPE.number,
      }));
  }, [feedbackScoresData?.scores]);

  const dynamicSpanScoresColumns = useMemo(() => {
    if (type !== TRACE_DATA_TYPE.traces) return [];
    return (spanFeedbackScoresData?.scores?.slice() ?? [])
      .sort((c1, c2) => c1.name.localeCompare(c2.name))
      .map<DynamicColumn>((c) => ({
        id: `${COLUMN_SPAN_FEEDBACK_SCORES_ID}.${c.name}`,
        label: formatSpanScoreLabel(c.name),
        columnType: COLUMN_TYPE.number,
      }));
  }, [spanFeedbackScoresData?.scores, type]);

  const dynamicColumnsIds = useMemo(
    () => [
      ...dynamicScoresColumns.map((c) => c.id),
      ...dynamicSpanScoresColumns.map((c) => c.id),
    ],
    [dynamicScoresColumns, dynamicSpanScoresColumns],
  );

  useDynamicColumnsCache({
    dynamicColumnsKey: DYNAMIC_COLUMNS_KEY,
    dynamicColumnsIds,
    setSelectedColumns,
  });

  const scoresColumnsData = useMemo(() => {
    // Always include "User feedback" column, even if it has no data
    const userFeedbackColumn: DynamicColumn = {
      id: USER_FEEDBACK_COLUMN_ID,
      label: USER_FEEDBACK_NAME,
      columnType: COLUMN_TYPE.number,
    };

    // Filter out "User feedback" from dynamic columns to avoid duplicates
    const otherDynamicColumns = dynamicScoresColumns.filter(
      (col) => col.id !== USER_FEEDBACK_COLUMN_ID,
    );

    const feedbackScoresColumns = [
      userFeedbackColumn,
      ...otherDynamicColumns,
    ].map(
      ({ label, id, columnType }) =>
        ({
          id,
          label,
          type: columnType,
          header: FeedbackScoreHeader as never,
          cell: FeedbackScoreCell as never,
          accessorFn: (row) =>
            row.feedback_scores?.find((f) => f.name === label),
          statisticKey: `${COLUMN_FEEDBACK_SCORES_ID}.${label}`,
        }) as ColumnData<BaseTraceData>,
    );

    // Add span feedback scores columns only for traces
    if (type === TRACE_DATA_TYPE.traces) {
      const spanFeedbackScoresColumns = dynamicSpanScoresColumns.map(
        ({ label, id, columnType }) => {
          // Extract the score name without the "(span)" suffix for matching
          const scoreName = parseSpanScoreName(label);
          return {
            id,
            label,
            type: columnType,
            header: FeedbackScoreHeader as never,
            cell: FeedbackScoreCell as never,
            accessorFn: (row) =>
              (row as Trace).span_feedback_scores?.find(
                (f) => f.name === scoreName,
              ),
            statisticKey: `${COLUMN_SPAN_FEEDBACK_SCORES_ID}.${scoreName}`,
          } as ColumnData<BaseTraceData>;
        },
      );
      return [...feedbackScoresColumns, ...spanFeedbackScoresColumns];
    }

    return feedbackScoresColumns;
  }, [dynamicScoresColumns, dynamicSpanScoresColumns, type]);

  const selectedRows: Array<Trace | Span> = useMemo(() => {
    return rows.filter((row) => rowSelection[row.id]);
  }, [rowSelection, rows]);

  const getDataForExport = useCallback(async (): Promise<
    Array<Trace | Span>
  > => {
    const result = await refetchExportData();

    if (result.error) {
      throw result.error;
    }

    if (!result.data?.content) {
      throw new Error("Failed to fetch data");
    }

    const allRows = result.data.content;
    const selectedIds = Object.keys(rowSelection);

    return allRows.filter((row) => selectedIds.includes(row.id));
  }, [refetchExportData, rowSelection]);

  const handleRowClick = useCallback(
    (row?: Trace | Span, lastSection?: DetailsActionSectionValue) => {
      if (!row) return;
      if (type === TRACE_DATA_TYPE.traces) {
        setTraceId((state) => (row.id === state ? "" : row.id));
        setSpanId("");
      } else {
        setTraceId((row as Span).trace_id);
        setSpanId((state) => (row.id === state ? "" : row.id));
      }

      if (lastSection) {
        setLastSection(lastSection);
      }
    },
    [setTraceId, setSpanId, type, setLastSection],
  );

  const meta = useMemo(
    () => ({
      onCommentsReply: (row?: Trace | Span) => {
        handleRowClick(row, DetailsActionSection.Comments);
      },
      enableUserFeedbackEditing: true,
    }),
    [handleRowClick],
  );

  const handleThreadIdClick = useCallback(
    (row?: Trace) => {
      if (!row) return;
      setThreadId(row.thread_id);
      setTraceId(row.id);
    },
    [setThreadId, setTraceId],
  );

  const columnData = useMemo(() => {
    return [
      ...SHARED_COLUMNS,
      ...(type === TRACE_DATA_TYPE.traces
        ? [
            {
              id: "span_count",
              label: "Span count",
              type: COLUMN_TYPE.number,
              accessorFn: (row: BaseTraceData) => get(row, "span_count", "-"),
            },
            {
              id: "llm_span_count",
              label: "LLM calls count",
              type: COLUMN_TYPE.number,
              accessorFn: (row: BaseTraceData) =>
                get(row, "llm_span_count", "-"),
            },
            {
              id: "thread_id",
              label: "Thread ID",
              type: COLUMN_TYPE.string,
              cell: LinkCell as never,
              customMeta: {
                callback: handleThreadIdClick,
                asId: true,
              },
              explainer: EXPLAINERS_MAP[EXPLAINER_ID.what_are_threads],
            },
          ]
        : []),
      ...(type === TRACE_DATA_TYPE.spans
        ? [
            {
              id: "type",
              label: "Type",
              type: COLUMN_TYPE.category,
              cell: SpanTypeCell as never,
            },
          ]
        : []),
      {
        id: "error_info",
        label: "Errors",
        statisticKey: "error_count",
        type: COLUMN_TYPE.errors,
        cell: ErrorCell as never,
      },
      {
        id: "created_by",
        label: "Created by",
        type: COLUMN_TYPE.string,
      },
      {
        id: COLUMN_COMMENTS_ID,
        label: "Comments",
        type: COLUMN_TYPE.string,
        cell: CommentsCell as never,
      },
      ...(isGuardrailsEnabled
        ? [
            {
              id: COLUMN_GUARDRAILS_ID,
              label: "Guardrails",
              statisticKey: COLUMN_GUARDRAIL_STATISTIC_ID,
              type: COLUMN_TYPE.category,
              iconType: "guardrails" as HeaderIconType,
              accessorFn: (row: BaseTraceData) =>
                row.guardrails_validations || [],
              cell: GuardrailsCell as never,
              statisticDataFormater: (value: number) => `${value} failed`,
            },
          ]
        : []),
    ];
  }, [type, handleThreadIdClick, isGuardrailsEnabled]);

  const filtersColumnData = useMemo(() => {
    return [
      {
        id: COLUMN_ID_ID,
        label: "ID",
        type: COLUMN_TYPE.string,
      },
      ...SHARED_COLUMNS,
      ...(type === TRACE_DATA_TYPE.traces
        ? [
            {
              id: "thread_id",
              label: "Thread ID",
              type: COLUMN_TYPE.string,
            },
            {
              id: "llm_span_count",
              label: "LLM calls count",
              type: COLUMN_TYPE.number,
            },
          ]
        : []),
      ...(type === TRACE_DATA_TYPE.spans
        ? [
            {
              id: "type",
              label: "Type",
              type: COLUMN_TYPE.category,
            },
          ]
        : []),
      {
        id: "error_info",
        label: "Errors",
        type: COLUMN_TYPE.errors,
      },
      {
        id: COLUMN_FEEDBACK_SCORES_ID,
        label: "Feedback scores",
        type: COLUMN_TYPE.numberDictionary,
      },
      ...(type === TRACE_DATA_TYPE.traces
        ? [
            {
              id: COLUMN_SPAN_FEEDBACK_SCORES_ID,
              label: "Span feedback scores",
              type: COLUMN_TYPE.numberDictionary,
            },
          ]
        : []),
      {
        id: COLUMN_CUSTOM_ID,
        label: "Custom filter",
        type: COLUMN_TYPE.dictionary,
      },
      ...(isGuardrailsEnabled
        ? [
            {
              id: COLUMN_GUARDRAILS_ID,
              label: "Guardrails",
              type: COLUMN_TYPE.category,
            },
          ]
        : []),
    ];
  }, [type, isGuardrailsEnabled]);

  const columns = useMemo(() => {
    return [
      generateSelectColumDef<Trace | Span>(),
      mapColumnDataFields<BaseTraceData, Span | Trace>({
        id: COLUMN_ID_ID,
        label: "ID",
        type: COLUMN_TYPE.string,
        cell: LinkCell as never,
        customMeta: {
          callback: handleRowClick,
          asId: true,
        },
        sortable: isColumnSortable(COLUMN_ID_ID, sortableBy),
      }),
      ...convertColumnDataToColumn<BaseTraceData, Span | Trace>(columnData, {
        columnsOrder,
        selectedColumns,
        sortableColumns: sortableBy,
      }),
      ...convertColumnDataToColumn<BaseTraceData, Span | Trace>(
        scoresColumnsData,
        {
          columnsOrder: scoresColumnsOrder,
          selectedColumns,
          sortableColumns: sortableBy,
        },
      ),
    ];
  }, [
    handleRowClick,
    sortableBy,
    columnData,
    columnsOrder,
    selectedColumns,
    scoresColumnsData,
    scoresColumnsOrder,
  ]);

  const columnsToExport = useMemo(() => {
    return columns
      .map((c) => get(c, "accessorKey", ""))
      .filter((c) =>
        c === COLUMN_SELECT_ID
          ? false
          : selectedColumns.includes(c) ||
            (DEFAULT_TRACES_COLUMN_PINNING.left || []).includes(c),
      );
  }, [columns, selectedColumns]);

  const activeRowId = type === TRACE_DATA_TYPE.traces ? traceId : spanId;
  const rowIndex = findIndex(rows, (row) => activeRowId === row.id);

  const hasNext = rowIndex >= 0 ? rowIndex < rows.length - 1 : false;
  const hasPrevious = rowIndex >= 0 ? rowIndex > 0 : false;

  const handleRowChange = useCallback(
    (shift: number) => handleRowClick(rows[rowIndex + shift]),
    [handleRowClick, rowIndex, rows],
  );

  const handleClose = useCallback(() => {
    setThreadId("");
    setTraceId("");
    setSpanId("");
  }, [setSpanId, setTraceId, setThreadId]);

  const sortConfig = useMemo(
    () => ({
      enabled: true,
      sorting: sortedColumns,
      setSorting: setSortedColumns,
    }),
    [setSortedColumns, sortedColumns],
  );

  const resizeConfig = useMemo(
    () => ({
      enabled: true,
      columnSizing: columnsWidth,
      onColumnResize: setColumnsWidth,
    }),
    [columnsWidth, setColumnsWidth],
  );

  const columnSections = useMemo(() => {
    return [
      {
        title: "Feedback scores",
        columns: scoresColumnsData,
        order: scoresColumnsOrder,
        onOrderChange: setScoresColumnsOrder,
      },
    ];
  }, [scoresColumnsData, scoresColumnsOrder, setScoresColumnsOrder]);

  if (isPending || isFeedbackScoresPending || isSpanFeedbackScoresPending) {
    return <Loader />;
  }

  return (
    <>
      <PageBodyStickyContainer direction="horizontal" limitWidth>
        <ExplainerCallout
          className="mb-4"
          {...(type === TRACE_DATA_TYPE.traces
            ? EXPLAINERS_MAP[EXPLAINER_ID.what_are_traces]
            : EXPLAINERS_MAP[EXPLAINER_ID.what_are_spans])}
        />
      </PageBodyStickyContainer>
      <PageBodyStickyContainer
        className="-mt-4 flex flex-wrap items-center justify-between gap-x-8 gap-y-2 py-4"
        direction="bidirectional"
        limitWidth
      >
        <div className="flex items-center gap-2">
          <SearchInput
            searchText={search as string}
            setSearchText={setSearch}
            placeholder="Search by ID"
            className="w-[320px]"
            dimension="sm"
          ></SearchInput>
          <FiltersButton
            columns={filtersColumnData}
            config={filtersConfig as never}
            filters={filters}
            onChange={setFilters}
          />
        </div>
        <div className="flex items-center gap-2">
          <TracesActionsPanel
            projectId={projectId}
            projectName={projectName}
            getDataForExport={getDataForExport}
            selectedRows={selectedRows}
            columnsToExport={columnsToExport}
            type={type as TRACE_DATA_TYPE}
          />
          <Separator orientation="vertical" className="mx-2 h-4" />
          <MetricDateRangeSelect
            value={dateRange}
            onChangeValue={handleDateRangeChange}
            minDate={minDate}
            maxDate={maxDate}
          />
          <TooltipWrapper
            content={`Refresh ${
              type === TRACE_DATA_TYPE.traces ? "traces" : "spans"
            } list`}
          >
            <Button
              variant="outline"
              size="icon-sm"
              className="shrink-0"
              onClick={() => {
                refetch();
                refetchStatistic();
              }}
            >
              <RotateCw />
            </Button>
          </TooltipWrapper>
          <DataTableRowHeightSelector
            type={height as ROW_HEIGHT}
            setType={setHeight}
          />
          <ColumnsButton
            columns={columnData}
            selectedColumns={selectedColumns}
            onSelectionChange={setSelectedColumns}
            order={columnsOrder}
            onOrderChange={setColumnsOrder}
            sections={columnSections}
          ></ColumnsButton>
        </div>
      </PageBodyStickyContainer>

      {noData && rows.length === 0 && page === 1 ? (
        <NoTracesPage />
      ) : (
        <>
          <DataTable
            columns={columns}
            columnsStatistic={columnsStatistic}
            data={rows}
            onRowClick={handleRowClick}
            activeRowId={activeRowId ?? ""}
            sortConfig={sortConfig}
            resizeConfig={resizeConfig}
            selectionConfig={{
              rowSelection,
              setRowSelection,
            }}
            getRowId={getRowId}
            rowHeight={height as ROW_HEIGHT}
            columnPinning={DEFAULT_TRACES_COLUMN_PINNING}
            noData={<DataTableNoData title={noDataText} />}
            TableWrapper={PageBodyStickyTableWrapper}
            stickyHeader
            meta={meta}
          />
          <PageBodyStickyContainer
            className="py-4"
            direction="horizontal"
            limitWidth
          >
            <DataTablePagination
              page={page as number}
              pageChange={setPage}
              size={size as number}
              sizeChange={setSize}
              total={data?.total ?? 0}
              supportsTruncation
              truncationEnabled={truncationEnabled}
            />
          </PageBodyStickyContainer>
        </>
      )}
      <TraceDetailsPanel
        projectId={projectId}
        traceId={traceId!}
        spanId={spanId!}
        setSpanId={setSpanId}
        setThreadId={setThreadId}
        hasPreviousRow={hasPrevious}
        hasNextRow={hasNext}
        open={Boolean(traceId) && !threadId}
        onClose={handleClose}
        onRowChange={handleRowChange}
      />
      <ThreadDetailsPanel
        projectId={projectId}
        projectName={projectName}
        traceId={traceId!}
        setTraceId={setTraceId}
        threadId={threadId!}
        open={Boolean(threadId)}
        onClose={handleClose}
      />
    </>
  );
};

export default TracesSpansTab;
